package com.mall.util.tree.rb;

import lombok.var;

public class RBTreeImpl<K extends Comparable<K>, V> implements RBTree<K, V> {
    /**
     * 树里面都需要一个根结点
     */
    private RBNode root;

    @Override
    public void insertNode(RBNode node) {
        if(this.root != null && this.root.getLeft() == null) {
            this.root.setLeft(node);
            node.setParent(this.root);
        }else {
            RBNode parent = findParentNode(node);

            /* 根据当前父节点，判断是否在父结点左边 */
            if (parent != null) {
                node.setParent(parent);

                int compareTo = node.getKey().compareTo(parent.getKey());

                if(compareTo == 0) {
                    parent.setValue(node.getValue());
                    return;
                }

                if (compareTo > 0) {
                    parent.setRight(node);
                } else if (compareTo < 0) {
                    parent.setLeft(node);
                }
            } else {
                this.root = node;
                this.root.setParent(null);
            }

            /* 插入之后，启动自平衡 */
            selfBalancing(node);
        }
    }

    @Override
    public void deleteNode(RBNode node) {

    }

    @Override
    public RBNode searchNode(K k) {
        return searchNode(new RBNode(k, null));
    }

    @Override
    public RBNode searchNode(RBNode node) {
        var p = this.root;

        while (p != null) {
            int compare = node.getKey().compareTo(p.getKey());
            if (compare > 0) {
                p = p.getRight();
            } else if (compare < 0) {
                p = p.getLeft();
            } else {
                return p;
            }
        }

        return null;
    }

    /**
     * 打印红黑树
     */
    @Override
    public void printTree() {
        var queue = new java.util.LinkedList<RBNode>();
        var queue2 = new java.util.LinkedList<RBNode>();
        if (root == null) {
            return;
        }
        queue.add(root);
        boolean firstQueue = true;

        while (!queue.isEmpty() || !queue2.isEmpty()) {
            var q = firstQueue ? queue : queue2;
            RBNode n = q.poll();

            if (n != null) {
                String pos = n.getParent() == null ? "" : (n == n.getParent().getLeft()
                        ? " LE" : " RI");
                String pstr = n.getParent() == null ? "" : n.getParent().toString();
                String cstr = RBColor.RED.equals(n.getColor()) ? "R" : "B";
                cstr = n.getParent() == null ? cstr : cstr + " ";
                System.out.print(n + "(" + (cstr) + pstr + (pos) + ")" + "\t");
                if (n.getLeft() != null) {
                    (firstQueue ? queue2 : queue).add(n.getLeft());
                }
                if (n.getRight() != null) {
                    (firstQueue ? queue2 : queue).add(n.getRight());
                }
            } else {
                System.out.println();
                firstQueue = !firstQueue;
            }
        }
    }

    /**
     * 结点右旋
     * p                 p
     * |                 |
     * y(旋转)             x
     * /   \     =>      /   \
     * x     ry         lx     y
     * /  \                    / \
     * lx   ly                ly   ry
     */
    private void rightRotate(RBNode y) {
        var x = y.getLeft();
        if(x == null){
            throw new java.lang.IllegalStateException("left node is null");
        }

        y.setLeft(x.getRight());
        if (x.getRight() != null) {
            x.getRight().setParent(y);
        }

        var p = y.getParent();
        if (p != null) {
            x.setParent(p);
            if (y == p.getRight()) {
                p.setRight(x);
            } else {
                p.setLeft(x);
            }
        } else {
            this.root = y;
            this.root.setParent(null);
        }

        x.setRight(y);
        y.setParent(x);
    }

    /**
     * 结点左旋
     * p                                 p
     * |                                 |
     * x               ===》》            y
     * /  \                              / \
     * lx    y                         x   ry
     * / \                          / \
     * ly   ry                       lx  ly
     */
    private void leftRotate(RBNode x) {
        var y = x.getRight();
        if(y == null) {
            throw new java.lang.IllegalStateException("right node is null");
        }
        x.setRight(y.getLeft());
        if (y.getLeft() != null) {
            y.getLeft().setParent(x);
        }

        var p = x.getParent();
        if (p != null) {
            y.setParent(p);
            if (x == p.getLeft()) {
                p.setLeft(x);
            } else {
                p.setRight(x);
            }
        } else {
            this.root = x;
            this.root.setParent(null);
        }

        x.setParent(y);
        y.setLeft(x);
    }

    /**
     * 修复平衡树
     * gp
     * /  \
     * p    uncle
     * |
     * rbNode
     */
    private void selfBalancing(RBNode rbNode) {
        var p = rbNode.getParent();

        while (p != null && RBColor.RED.equals(p.getColor())) {
            var gp = p.getParent();
            var uncle = gp.getRight();

            if (uncle != null && RBColor.RED.equals(uncle.getColor())) {
                p.setColor(RBColor.BLACK);
                uncle.setColor(RBColor.BLACK);
                gp.setColor(RBColor.RED);

                /* 以爷爷结点进入下一轮处理 */
                rbNode = gp;
            }else {
                if(p == gp.getLeft()) {
                    if(rbNode == p.getRight()) {
                        leftRotate(p);

                        rbNode = p;
                        p = rbNode.getParent();
                    }

                    p.setColor(RBColor.BLACK);
                    gp.setColor(RBColor.RED);

                    rightRotate(gp);
                }else if(p == gp.getRight()){
                    if(rbNode == p.getLeft()) {
                        rightRotate(p);

                        rbNode = p;
                        p = rbNode.getParent();
                    }

                    p.setColor(RBColor.BLACK);
                    gp.setColor(RBColor.RED);

                    leftRotate(gp);
                }
            }
        }

        this.root.setColor(RBColor.BLACK);
    }

    /**
     * 找寻到当前结点的父亲结点
     */
    private RBNode findParentNode(RBNode rbNode) {
        var p = this.root;
        var child = p;

        while (child != null) {
            int cmp = child.getKey().compareTo(rbNode.getKey());
            if(cmp == 0) {
                return child;
            }

            if(cmp > 0) {
                p = child;
                child = child.getLeft();
            }else {
                p = child;
                child = child.getRight();
            }
        }
        return p;
    }

    /**
     * 树结点的中序遍历
     */
    private void minPrintf() {

    }
}
